/*
 * Copyright (c) 2024 endless-sky
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "mtd_manager.h"
#include "utils_log.h"
#include <fcntl.h>
#include <unistd.h>
#include <cerrno>
#include <algorithm>

namespace El {
namespace Upgrade {

MtdManager::MtdManager() : mtdLib_(nullptr) {}

MtdManager::~MtdManager()
{
    CleanupMtdResources();
}

bool MtdManager::InitializeMtdLib()
{
    EL_INFO("Initializing MTD library");

    // Initialize libmtd
    mtdLib_ = libmtd_open();
    if (!mtdLib_) {
        lastError_ = "Failed to initialize libmtd";
        EL_ERROR("{}", lastError_);
        return false;
    }

    return true;
}

bool MtdManager::ConfigurePartitions(const std::vector<MtdPartitionMap> &partitions)
{
    EL_INFO("Configuring MTD partitions");

    partitions_ = partitions;

    // 获取每个分区的详细信息并验证设备可访问性
    bool anyDeviceAccessible = false;
    for (auto &partition : partitions_) {
        MtdDevice mtd;
        mtd.fd = -1; // Initialize file descriptor
        if (OpenMtdDevice(partition.device, mtd)) {
            // 从实际设备获取所有参数
            const_cast<MtdPartitionMap &>(partition).size = mtd.dev.size;
            const_cast<MtdPartitionMap &>(partition).erasesize = mtd.dev.eb_size;
            const_cast<MtdPartitionMap &>(partition).min_io_size = mtd.dev.min_io_size;
            CloseMtdDevice(mtd);
            anyDeviceAccessible = true;
            EL_INFO("MTD partition {} ({}): size={}KB, erasesize={}KB, min_io_size={}B", partition.name,
                     partition.device, partition.size / 1024, partition.erasesize / 1024, partition.min_io_size);
        } else {
            lastError_ = "Failed to access MTD device: " + partition.device;
            EL_ERROR("{}", lastError_);
            return false;
        }
    }

    if (!anyDeviceAccessible) {
        lastError_ = "No MTD devices are accessible";
        EL_ERROR("{}", lastError_);
        return false;
    }

    EL_INFO("Successfully configured {} MTD partitions", partitions_.size());
    return true;
}

void MtdManager::CleanupMtdResources()
{
    if (mtdLib_) {
        libmtd_close(mtdLib_);
        mtdLib_ = nullptr;
    }
    partitions_.clear();
}

bool MtdManager::OpenMtdDevice(const std::string &device, MtdDevice &mtd)
{
    // 首先检查设备文件是否存在
    if (access(device.c_str(), F_OK) != 0) {
        EL_WARN("MTD device does not exist: {}", device);
        return false;
    }

    // 尝试打开设备，使用非阻塞模式避免卡死
    mtd.fd = open(device.c_str(), O_RDWR | O_NONBLOCK);
    if (mtd.fd < 0) {
        EL_WARN("Failed to open MTD device: {} (errno: {})", device, errno);
        return false;
    }

    // 设置回阻塞模式进行后续操作
    int flags = fcntl(mtd.fd, F_GETFL, 0);
    if (flags == -1 || fcntl(mtd.fd, F_SETFL, flags & ~O_NONBLOCK) == -1) {
        EL_WARN("Failed to set MTD device to blocking mode: {}", device);
        close(mtd.fd);
        return false;
    }

    mtd.mtd = mtdLib_;
    if (mtd_get_dev_info(mtdLib_, device.c_str(), &mtd.dev) < 0) {
        EL_WARN("Failed to get MTD device info: {}", device);
        close(mtd.fd);
        return false;
    }

    return true;
}

void MtdManager::CloseMtdDevice(MtdDevice &mtd)
{
    if (mtd.fd >= 0) {
        close(mtd.fd);
        mtd.fd = -1;
    }
}

bool MtdManager::EraseMtdDevice(MtdDevice &mtd, uint64_t offset, uint64_t size)
{
    // 按擦除块大小对齐
    uint64_t alignedOffset = (offset / mtd.dev.eb_size) * mtd.dev.eb_size;
    uint64_t alignedSize = ((size + mtd.dev.eb_size - 1) / mtd.dev.eb_size) * mtd.dev.eb_size;

    EL_DEBUG("Erasing MTD device: offset=0x{:x}, size=0x{:x}", alignedOffset, alignedSize);

    for (uint64_t addr = alignedOffset; addr < alignedOffset + alignedSize; addr += mtd.dev.eb_size) {
        if (mtd_erase(mtdLib_, &mtd.dev, mtd.fd, addr / mtd.dev.eb_size) < 0) {
            EL_ERROR("Failed to erase block at offset 0x{:x}", addr);
            return false;
        }
    }

    return true;
}

bool MtdManager::EraseMtdDeviceWithBadBlocks(MtdDevice &mtd, uint64_t offset, uint64_t size)
{
    // 按擦除块大小对齐
    uint64_t alignedOffset = (offset / mtd.dev.eb_size) * mtd.dev.eb_size;
    uint64_t alignedSize = ((size + mtd.dev.eb_size - 1) / mtd.dev.eb_size) * mtd.dev.eb_size;

    EL_DEBUG("Erasing MTD device with bad block handling: offset=0x{:x}, size=0x{:x}", alignedOffset, alignedSize);

    uint64_t badBlocksCount = 0;
    uint64_t totalBlocks = alignedSize / mtd.dev.eb_size;

    for (uint64_t addr = alignedOffset; addr < alignedOffset + alignedSize; addr += mtd.dev.eb_size) {
        uint64_t eraseBlockNum = addr / mtd.dev.eb_size;

        // 检查是否为坏块
        int badBlockResult = mtd_is_bad(&mtd.dev, mtd.fd, eraseBlockNum);
        if (badBlockResult < 0) {
            EL_ERROR("Failed to check bad block for block {} during erase", eraseBlockNum);
            return false;
        }

        if (badBlockResult == 1) {
            // 跳过坏块
            EL_WARN("Skipping bad block {} during erase at offset 0x{:x}", eraseBlockNum, addr);
            badBlocksCount++;
            continue;
        }

        // 尝试擦除好块
        const int maxRetries = 3;
        int retryCount = 0;
        bool eraseSuccess = false;

        while (retryCount < maxRetries && !eraseSuccess) {
            if (mtd_erase(mtdLib_, &mtd.dev, mtd.fd, eraseBlockNum) >= 0) {
                eraseSuccess = true;
            } else {
                retryCount++;
                if (retryCount < maxRetries) {
                    EL_WARN("Erase failed for block {}, retry {}/{}", eraseBlockNum, retryCount, maxRetries);
                } else {
                    EL_ERROR("Failed to erase block {} after {} retries", eraseBlockNum, maxRetries);
                    // 标记为坏块
                    if (mtd_mark_bad(&mtd.dev, mtd.fd, eraseBlockNum) < 0) {
                        EL_WARN("Failed to mark block {} as bad during erase", eraseBlockNum);
                    }
                    badBlocksCount++;
                }
            }
        }

        if (!eraseSuccess) {
            // 擦除失败，跳过这个块
            continue;
        }
    }

    EL_INFO("Erase completed: {} total blocks, {} bad blocks skipped", totalBlocks, badBlocksCount);
    return true;
}

bool MtdManager::WriteImageStream(const std::string &partitionName,
                                    const uint8_t *data,
                                    size_t length,
                                    bool isLastChunk,
                                    std::map<std::string, uint64_t> &writeOffsets,
                                    std::map<std::string, uint64_t> &writeLogicalOffsets,
                                    std::map<std::string, bool> &writeStarted,
                                    std::map<std::string, std::vector<uint8_t>> &writeBuffers)
{
    EL_DEBUG("Writing image stream: partition={}, length={}, last={}", partitionName, length, isLastChunk);

    auto *partition = GetPartition(partitionName);
    if (!partition) {
        lastError_ = "Partition not found: " + partitionName;
        EL_ERROR("{}", lastError_);
        return false;
    }

    // Open MTD device
    MtdDevice mtd;
    mtd.fd = -1; // Initialize file descriptor
    if (!OpenMtdDevice(partition->device, mtd)) {
        return false;
    }

    // Get current write offset
    uint64_t &offset = writeOffsets[partitionName];
    uint64_t &logicalOffset = writeLogicalOffsets[partitionName];

    // Erase partition on first write (with bad block handling)
    if (!writeStarted[partitionName]) {
        if (!EraseMtdDeviceWithBadBlocks(mtd, 0, partition->size)) {
            CloseMtdDevice(mtd);
            return false;
        }
        writeStarted[partitionName] = true;
        offset = 0;
        logicalOffset = 0;
    }

    // 使用缓冲区处理流式数据，但简化对齐逻辑
    std::vector<uint8_t> &writeBuffer = writeBuffers[partitionName];
    writeBuffer.insert(writeBuffer.end(), data, data + length);

    // 严格按照擦除块大小写入数据
    size_t eraseBlockSize = mtd.dev.eb_size;

    // 只有当缓冲区有足够数据时才写入完整的擦除块
    while (writeBuffer.size() >= eraseBlockSize) {
        // 计算当前擦除块号
        uint64_t eraseBlockNum = offset / eraseBlockSize;

        // 检查是否为坏块
        int badBlockResult = mtd_is_bad(&mtd.dev, mtd.fd, eraseBlockNum);
        if (badBlockResult < 0) {
            EL_ERROR("Failed to check bad block for block {}", eraseBlockNum);
            CloseMtdDevice(mtd);
            return false;
        }

        if (badBlockResult == 1) {
            // 跳过坏块
            EL_WARN("Skipping bad block {} at offset 0x{:x}", eraseBlockNum, offset);
            offset += eraseBlockSize;
            continue;
        }

        // 尝试写入，带重试机制
        const int maxRetries = 3;
        int retryCount = 0;
        bool writeSuccess = false;

        while (retryCount < maxRetries && !writeSuccess) {
            if (mtd_write(mtdLib_, &mtd.dev, mtd.fd, eraseBlockNum, 0, const_cast<uint8_t *>(writeBuffer.data()),
                          eraseBlockSize, nullptr, 0, 0) >= 0) {
                writeSuccess = true;
            } else {
                retryCount++;
                if (retryCount < maxRetries) {
                    EL_WARN("Write failed for block {}, retry {}/{}", eraseBlockNum, retryCount, maxRetries);
                    // 重新擦除块
                    if (mtd_erase(mtdLib_, &mtd.dev, mtd.fd, eraseBlockNum) < 0) {
                        EL_ERROR("Failed to re-erase block {} during retry", eraseBlockNum);
                        break;
                    }
                } else {
                    EL_ERROR("Failed to write erase block {} after {} retries", eraseBlockNum, maxRetries);
                    // 标记为坏块
                    if (mtd_mark_bad(&mtd.dev, mtd.fd, eraseBlockNum) < 0) {
                        EL_WARN("Failed to mark block {} as bad", eraseBlockNum);
                    }
                }
            }
        }

        if (!writeSuccess) {
            CloseMtdDevice(mtd);
            return false;
        }

        offset += eraseBlockSize;
        logicalOffset += eraseBlockSize;
        writeBuffer.erase(writeBuffer.begin(), writeBuffer.begin() + eraseBlockSize);

        EL_DEBUG("Wrote erase block {} ({} bytes), total offset: 0x{:x}", eraseBlockNum, eraseBlockSize, offset);
    }

    // 处理最后一个块
    if (isLastChunk && !writeBuffer.empty()) {
        // 填充到完整的擦除块大小
        writeBuffer.resize(eraseBlockSize, 0xFF);

        // 寻找下一个好的块来写入最后的数据
        const int maxBlockSearch = 10; // 最多搜索10个块
        int blocksSearched = 0;
        bool finalWriteSuccess = false;

        while (blocksSearched < maxBlockSearch && !finalWriteSuccess) {
            uint64_t eraseBlockNum = offset / eraseBlockSize;

            // 检查是否为坏块
            int badBlockResult = mtd_is_bad(&mtd.dev, mtd.fd, eraseBlockNum);
            if (badBlockResult < 0) {
                EL_ERROR("Failed to check bad block for final block {}", eraseBlockNum);
                break;
            }

            if (badBlockResult == 1) {
                EL_WARN("Skipping bad block {} for final write at offset 0x{:x}", eraseBlockNum, offset);
                offset += eraseBlockSize;
                blocksSearched++;
                continue;
            }

            // 尝试写入最后的数据
            const int maxRetries = 3;
            int retryCount = 0;

            while (retryCount < maxRetries && !finalWriteSuccess) {
                if (mtd_write(mtdLib_, &mtd.dev, mtd.fd, eraseBlockNum, 0, const_cast<uint8_t *>(writeBuffer.data()),
                              eraseBlockSize, nullptr, 0, 0) >= 0) {
                    finalWriteSuccess = true;
                } else {
                    retryCount++;
                    if (retryCount < maxRetries) {
                        EL_WARN("Final write failed for block {}, retry {}/{}", eraseBlockNum, retryCount, maxRetries);
                        // 重新擦除块
                        if (mtd_erase(mtdLib_, &mtd.dev, mtd.fd, eraseBlockNum) < 0) {
                            EL_ERROR("Failed to re-erase block {} during final retry", eraseBlockNum);
                            break;
                        }
                    } else {
                        EL_ERROR("Failed to write final erase block {} after {} retries", eraseBlockNum, maxRetries);
                        // 标记为坏块
                        if (mtd_mark_bad(&mtd.dev, mtd.fd, eraseBlockNum) < 0) {
                            EL_WARN("Failed to mark final block {} as bad", eraseBlockNum);
                        }
                    }
                }
            }

            if (finalWriteSuccess) {
                offset += eraseBlockSize;
                logicalOffset += eraseBlockSize;
                EL_DEBUG("Wrote final erase block {} ({} bytes), total offset: 0x{:x}", eraseBlockNum, eraseBlockSize,
                          offset);
            } else {
                offset += eraseBlockSize;
                blocksSearched++;
            }
        }

        if (!finalWriteSuccess) {
            EL_ERROR("Failed to write final data after searching {} blocks", maxBlockSearch);
            CloseMtdDevice(mtd);
            return false;
        }

        writeBuffer.clear();
    }

    CloseMtdDevice(mtd);
    return true;
}

MtdManager::MtdPartitionMap *MtdManager::GetPartition(const std::string &name)
{
    for (auto &partition : partitions_) {
        if (partition.name == name) {
            return &partition;
        }
    }
    return nullptr;
}

void MtdManager::SetLastError(const std::string &error)
{
    lastError_ = error;
}

std::string MtdManager::GetLastError() const
{
    return lastError_;
}

} // namespace Upgrade
} // namespace El
